from utils.utils import *
from config.config import *
import joblib
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import os

# ===================== 读取特征 =====================
signature_list = ['Collagen','Clinic','Nuclei','Region']
clinic_select_features = [ 'Tumour Size(cm)', 'Age(years)', 'ER','PR','HER2']

collagen_columns = pd.read_csv(opj(lasso_feature_path,str(split_random_state_list[0]),str(selectK_random_state_list[0]),f'Collagen.csv')).drop(columns=exclude_columns).columns
region_columns = pd.read_csv(opj(lasso_feature_path,str(split_random_state_list[0]),str(selectK_random_state_list[0]),f'Region.csv')).drop(columns=exclude_columns).columns
nuclei_columns = pd.read_csv(opj(lasso_feature_path,str(split_random_state_list[0]),str(selectK_random_state_list[0]),f'Roi.csv')).drop(columns=exclude_columns).columns
clinic_columns = pd.read_csv(clinic_data_path)[clinic_select_features].columns

# collagen_columns = pd.Series(collagen_columns).str.extract(r'\((.*?)\)')[0].tolist()
# region_columns = pd.Series(region_columns).str.extract(r'\((.*?)\)')[0].tolist()
# nuclei_columns = pd.Series(nuclei_columns).str.extract(r'\((.*?)\)')[0].tolist()
# clinic_columns = pd.Series(clinic_columns).str.replace(r'\s*\(.*\)', '', regex=True).tolist()

# ===================== 加载模型 =====================
collagen_model = joblib.load(opj(signature_model_weight_path,'Collagen','model.pkl')).get_model()
region_model = joblib.load(opj(signature_model_weight_path,'Region','model.pkl')).get_model()
nuclei_model = joblib.load(opj(signature_model_weight_path,'Nuclei','model.pkl')).get_model()
clinic_model = joblib.load(opj(signature_model_weight_path,'Clinic','model.pkl')).get_model()

# ===================== 构建 DataFrame =====================
coef_dfs = []

collagen_coef = pd.DataFrame({'Feature': collagen_columns, 'Coefficient': collagen_model.coef_[0], 'Group': 'Collagen'})
region_coef = pd.DataFrame({'Feature': region_columns, 'Coefficient': region_model.coef_[0], 'Group': 'Region'})
nuclei_coef = pd.DataFrame({'Feature': nuclei_columns, 'Coefficient': nuclei_model.coef_[0], 'Group': 'Nuclei'})
clinic_coef = pd.DataFrame({'Feature': clinic_columns, 'Coefficient': clinic_model.coef_[0], 'Group': 'Clinic'})

coef_dfs = [collagen_coef, region_coef, nuclei_coef, clinic_coef]
md(signature_construction_model_visulation_result_path)

# 保存整体数据
coef_df_all = pd.concat(coef_dfs, ignore_index=True)
coef_df_all.to_csv(opj(signature_construction_model_visulation_result_path,'all_coef.csv'), index=False)

# ===================== 绘制热力图 =====================


heatmap_path = opj(signature_construction_model_visulation_result_path,'heatmap')
md(heatmap_path)
# 颜色映射



for df in coef_dfs:
    sig_name = df['Group'].iloc[0]
    df_plot = df.copy()
    df_plot = df_plot.set_index('Feature')
    plt.figure(figsize=(max(14, len(df_plot)*0.3), 8))
    max_abs = max(abs(df_plot['Coefficient'].min()), abs(df_plot['Coefficient'].max()))
    sns.heatmap(df_plot[['Coefficient']], annot=True, cmap='RdBu_r', center=0, linewidths=0.5,vmin=-max_abs, vmax=max_abs)
    plt.title(f"{sig_name} Signature Sub-feature Coefficients")
    plt.ylabel("Sub-feature")
    plt.yticks(rotation=0) 
    plt.tight_layout()
    plt.savefig(opj(heatmap_path, f"{sig_name}.png"), dpi=300)
    plt.show()
